import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import time
import matplotlib.dates as mdates

# 读取CSV文件
# file_path = 'logs\\backtest_price.csv'
file_path = 'logs/predict/price.csv'
df = pd.read_csv(file_path)

# 转换时间戳为datetime对象（毫秒级时间戳），并指定时区为UTC
df['datetime'] = pd.to_datetime(df['timestamp'], unit='ms', utc=True)
df['datetime'] = df['datetime'].dt.tz_convert('Asia/Shanghai')

# 按分钟取每分钟最后一条记录（最近数据）
df_1min = df.groupby(df['datetime'].dt.floor('T')).first()
# 去掉 tz 信息，方便后续处理
df_1min.index = df_1min.index.tz_convert('Asia/Shanghai').tz_localize(None)

# 记住原始（未删除中午前后的）索引用于寻找边界
orig_idx = df_1min.index

# 删除 11:30（含）到 13:00（不含）
midday = df_1min.between_time('11:30', '12:59:59.999999')
df_plot = df_1min.drop(midday.index)

if df_plot.empty:
    raise SystemExit("过滤后无绘图数据，请检查原始时间范围。")

# 计算真正的分时均线（成交量加权平均价）
# 分时均价 = 累计成交金额 ÷ 累计成交股数
df_plot['cum_total_trade_value'] = df_plot['total_trade_value'].cumsum()
df_plot['cum_total_trade_volume'] = df_plot['total_trade_volume'].cumsum()
df_plot['time_weighted_avg'] = df_plot['cum_total_trade_value'] / df_plot['cum_total_trade_volume']

# 用连续整数作为 x 值（压缩中午段）
x = np.arange(len(df_plot))
y_price = df_plot['price'].values
y_time_avg = df_plot['time_weighted_avg'].values

# 准备刻度：每30分钟显示一个刻度
ticks = []
labels = []

# 遍历所有数据点，选择每30分钟的点作为刻度
for i, dt in enumerate(df_plot.index):
    # 如果是整点或30分钟，且不是重复的时间
    if dt.minute == 0 or dt.minute == 30:
        ticks.append(i)
        labels.append(dt.strftime('%H:%M'))

# 确保第一个和最后一个点都有刻度
if len(df_plot) > 0:
    if 0 not in ticks:
        ticks.insert(0, 0)
        labels.insert(0, df_plot.index[0].strftime('%H:%M'))
    
    if (len(df_plot) - 1) not in ticks:
        ticks.append(len(df_plot) - 1)
        labels.append(df_plot.index[-1].strftime('%H:%M'))

# 绘图（压缩中午段）
fig, ax = plt.subplots(figsize=(15, 6))

# 绘制两条线：价格线、分时均线
ax.plot(x, y_price, linewidth=2, marker='o', markersize=3, alpha=0.8, 
        label='Price', color='#1f77b4')
ax.plot(x, y_time_avg, linewidth=2.5, label='Time Weighted Average', 
        color='#ff7f0e', alpha=0.9)

# 设刻度（整数位置），并显示对应的真实时间标签
ax.set_xticks(ticks)
ax.set_xticklabels(labels, rotation=45)

ax.set_title("Stock Software Style Chart with Time Weighted Average", fontsize=14, fontweight='bold')
ax.set_xlabel("Time")
ax.set_ylabel("Price")
ax.grid(True, alpha=0.3)
ax.legend()

# 在11:30断开处添加垂直线标注
morning_end_idx = None
afternoon_start_idx = None

# 找到上午结束和下午开始的位置
for i, dt in enumerate(df_plot.index):
    if dt.time() < time(11, 30):
        morning_end_idx = i
    elif dt.time() >= time(13, 0) and afternoon_start_idx is None:
        afternoon_start_idx = i
        break

# 在断开处添加垂直线
if morning_end_idx is not None and afternoon_start_idx is not None:
    # 在断开处添加垂直线
    ax.axvline(x=morning_end_idx + 0.5, color='red', linestyle='--', alpha=0.5, linewidth=1)
    ax.axvline(x=afternoon_start_idx - 0.5, color='red', linestyle='--', alpha=0.5, linewidth=1)
    
    # 添加文本标注
    ax.text(morning_end_idx + 0.5, np.max(y_price), 'Lunch Break', 
            rotation=90, verticalalignment='top', horizontalalignment='right',
            fontsize=10, color='red', alpha=0.7)

plt.tight_layout()
plt.show()

# 打印处理信息
print("数据处理信息:")
print(f"原始数据点数: {len(df)}")
print(f"分钟级数据点数: {len(df_1min)}")
print(f"过滤后数据点数: {len(df_plot)}")
print(f"时间范围: {df_plot.index[0].strftime('%H:%M')} 到 {df_plot.index[-1].strftime('%H:%M')}")
print(f"删除的中午时段: 11:30 - 13:00 ({len(midday)} 个数据点)")

if morning_end_idx is not None and afternoon_start_idx is not None:
    print(f"上午结束时间: {df_plot.index[morning_end_idx].strftime('%H:%M')}")
    print(f"下午开始时间: {df_plot.index[afternoon_start_idx].strftime('%H:%M')}")

# 显示详细的统计信息
print(f"\n价格统计信息:")
print(f"价格范围: {df_plot['price'].min():.4f} - {df_plot['price'].max():.4f}")
print(f"最终价格: {df_plot['price'].iloc[-1]:.4f}")

print(f"\n分时均线统计信息:")
print(f"分时均线范围: {df_plot['time_weighted_avg'].min():.4f} - {df_plot['time_weighted_avg'].max():.4f}")
print(f"最终分时均线值: {df_plot['time_weighted_avg'].iloc[-1]:.4f}")

print(f"\n成交量统计信息:")
print(f"总成交量: {df_plot['total_trade_volume'].sum():.0f} 股")
print(f"总成交金额: {df_plot['total_trade_value'].sum():.2f} 元")

print(f"\n价格与分时均线对比:")
price_final = df_plot['price'].iloc[-1]
time_avg_final = df_plot['time_weighted_avg'].iloc[-1]
difference = price_final - time_avg_final
if difference > 0:
    print(f"当前价格高于分时均线: +{difference:.4f}")
elif difference < 0:
    print(f"当前价格低于分时均线: {difference:.4f}")
else:
    print(f"当前价格等于分时均线")